
// Definition for a binary tree node.
struct TreeNode
{
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
};
#include <unordered_map>
#include <vector>
using namespace std;
//中序遍历
class Solution
{
public:
    pair<bool, TreeNode *> helper(TreeNode *root, TreeNode *target)
    {
        if (!root)
            return {false, nullptr};
        // 先找左边
        auto left_res = helper(root->left, target);
        // 如果左边找到了但没有答案代表是左边最后一个节点 答案为根
        if (left_res.first)
            return left_res.second ? left_res : make_pair(true, root);
        // 左边没找到 如果根是target
        if (root == target)
        {
            // 答案为右边的最左
            if (root->right)
            {
                TreeNode *res = root->right;
                while (res && res->left)
                    res = res->left;
                return {true, res};
            }
            // 没有右边 代表找到target但得不到答案
            return {true, nullptr};
        }
        // 最后找右边
        return helper(root->right, target);
    }
    TreeNode *inorderSuccessor(TreeNode *root, TreeNode *p)
    {
        auto res = helper(root, p);
        return res.second;
    }
};
// 空间换时间
class Solution
{
    vector<TreeNode *> arr;
    unordered_map<TreeNode *, int> hash_map;

public:
    void build_index(TreeNode *root)
    {
        if (!root)
            return;
        build_index(root->left);
        arr.push_back(root);
        hash_map[root] = arr.size() - 1;
        build_index(root->right);
    }
    TreeNode *inorderSuccessor(TreeNode *root, TreeNode *p)
    {
        build_index(root);
        auto it = hash_map.find(p);
        if (it == hash_map.end() || it->second + 1 == arr.size())
            return nullptr;
        return arr[it->second + 1];
    }
};
